import tornado.ioloop
import tornado.web
from tornado.ioloop import IOLoop
import os
import torch
from torch import nn
import argparse
import json

import google.protobuf as pb
print(f'protobuf version: {pb.__version__}')
from YAudio2Bs_pb2 import Audio2BSData
import threading

from BSGenerator import BSGenerator,BSGeneratorV2
import base64
import numpy as np
import time
from global_config import *
from txt2json import *

from concurrent import futures
from google.cloud import pubsub_v1
from typing import Callable
import sys
from rocketmq import ClientConfiguration, Credentials, Message, Producer
from neurosync.BSGenerator import BSGenerator as BSGeneratorNS
lock = threading.Lock()
import subprocess

generatorV1 = None
generatorV2 = None
generatorNS = None
use_version = 'v1'
git_version = ""
def choose_generator(version):
    if version == 'v1':
        generator = generatorV1
    elif version == 'v2':        
        generator = generatorV2
    elif version == 'ns':        
        generator = generatorNS
    return generator
      
use_google = False
google_publisher = None
google_topic_path = None
use_envelope = False

use_aliyun = False
aliyun_publisher = None
aliyun_topic_path = None
       
def jsonToProto(jdata,fileId,version):
    pdata = Audio2BSData()
    for frame in jdata.frames:
        pframe = pdata.frames.add()
        for location in frame.location:
            pframe.location.append(location)
        for weight in frame.weight:
            pframe.weight.append(weight)
        pframe.time = frame.time
    if fileId is not None:
        pdata.fileId = fileId
    if version is not None:
        pdata.version = version
    s = pdata.SerializeToString()
    return s 


requestCount = 0

def publish(data):    
    if use_google:
        try:
            # When you publish a message, the client returns a future.
            future = google_publisher.publish(google_topic_path, data)
            print(f'publish to google topic_path: {google_topic_path}')        
            print('publish future result:' + future.result())
            
        except Exception as e:
            print('google publish exception')
    if use_aliyun:
        try:
            aliyun_publisher.startup()
            try:
                msg = Message()
                msg.topic = aliyun_topic_path
                msg.body = data
                msg.tag = "rocketmq-send-message"
                msg.keys = "send_sync"
                msg.add_property("send", "sync")
                res = aliyun_publisher.send(msg)
                print(f"{aliyun_publisher.__str__()} send message success. {res}")
                aliyun_publisher.shutdown()
                print(f"{aliyun_publisher.__str__()} shutdown.")
            except Exception as e:
                print(f"normal producer example raise exception: {e}")
                aliyun_publisher.shutdown()
        except Exception as e:
            print(f"{aliyun_publisher.__str__()} startup raise exception: {e}")
            aliyun_publisher.shutdown()
            
            

def get_git_commit_id():
    try:
        commit_id = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('ascii').strip()
        return commit_id
    except subprocess.CalledProcessError:
        return None
    except FileNotFoundError:
        return None  
     
def do_common(handler):
    print('start do post')       
    
    global requestCount
    requestCount += 1       
    
    try: 
        if use_envelope == False:
            json_object =  tornado.escape.json_decode(handler.request.body)     
            base64_str = json_object['WavData']
            
            fileId = None
            if 'fileId' in json_object:
                fileId = json_object['fileId']
                
                
            if 'Version' not in json_object:
                version = 'v1'
            else:
                version = json_object['Version']
            
            if 'UseProtobuf' not in json_object:
                useproto = False
            else:
                useproto = json_object['UseProtobuf']
            
            byte_data = base64.b64decode(base64_str)#将base64转换为二进制
        else:
            try:
                envelope = json.loads(handler.request.body)
            except json.JSONDecodeError:
                handler.set_status(400)
                handler.write("Invalid JSON")
                return

            if "message" not in envelope:
                handler.set_status(400)
                handler.write("Invalid Pub/Sub message")
                return

            msg = envelope["message"]
            data = msg.get("data")

            if data:
                # 第一步：Base64 解码
                decoded = base64.b64decode(data).decode("utf-8")
                print("收到原始 JSON 字符串：", decoded)

                # 第二步：解析为字典
                body = json.loads(decoded)
                print("解析后的内容：", body)

                # 第三步：取出各个字段
                wav_data_b64 = body.get("WavData")
                file_id_json = body.get("fileId")

                # fileId 也是 JSON 字符串，再转一次
                file_id = json.loads(file_id_json)

                print("音频 Base64 长度：", len(wav_data_b64))
                print("文件标识对象：", file_id)
                print("UseProtobuf:", body.get("UseProtobuf"))
                print("Version:", body.get("Version"))

                useproto = body.get("UseProtobuf")
                version = body.get("Version")
                # 如果需要恢复音频文件：
                # with open("output.wav", "wb") as f:
                #     f.write(base64.b64decode(wav_data_b64))
                byte_data = base64.b64decode(wav_data_b64)#将base64转换为二进制

                
       
        if version == 'ns':
            input = byte_data
        else: 
            tempPath = './temp.wav'
            
            file = open(tempPath, 'wb')
            try:
                file.write(byte_data)
            finally:
                file.close()
            input = tempPath
            
    
        generator = choose_generator(version)
        data = generator.computeBS(input)   
        if useproto:
            if use_google or use_aliyun:                   
            
                pbytes = jsonToProto(data,fileId,git_version)
                publish(pbytes)
                
                dictdata = dict()
                dictdata['fileId'] = fileId
                dictdata['version'] = git_version
                dictdata['ret'] = 'OK'                
                jstr = json.dumps(dictdata)   
                handler.set_status(200)
                handler.set_header('Content-Type', 'application/json')
                
                # handler.write(pbytes)
                handler.write(jstr)
            else:
                pbytes = jsonToProto(data,fileId,git_version)               
                handler.set_status(200)
                handler.set_header('Content-Type', 'application/octet-stream')
                print(f"protobuf send data len: {len(pbytes)}")
                handler.write(pbytes)
                
            
        else:
            dictdata = asdict(data)
            dictdata['fileId'] = fileId
            dictdata['version'] = git_version
            dictdata['ret'] = 'OK'
            jstr = json.dumps(dictdata)   
            handler.set_status(200)
            handler.set_header('Content-Type', 'application/json')
            data = jstr.encode()
            print(f"json send data len: {len(data)}")
            handler.write(data)
            
            publish(data)
        
        print(f'reqcount {requestCount} Post Sucess')
    except Exception as error:
        print(f'reqcount {requestCount} Post Connection Error : {error}')
        
        handler.set_status(500)
        handler.set_header('Content-Type', 'application/json')
        jstr = '{}'
        data = jstr.encode()
        handler.write(data)
    

class MainHandler(tornado.web.RequestHandler):
    def get(self):
        data={}
        data['ret'] = "ok"
        jstr = json.dumps(data)   
        self.write(jstr)
    
    def post(self):
        
        lock.acquire()
        do_common(self)        
        lock.release()
        
        
class GCPHandler(tornado.web.RequestHandler):
    def get(self):
        data={}
        data['ret'] = "ok"
        jstr = json.dumps(data)   
        self.write(jstr)
        
    def post(self):
        lock.acquire()
        do_common(self)        
        lock.release()

def make_app():
    return tornado.web.Application([
        (r"/", MainHandler),
        (r"/gcp", GCPHandler),        
    ])
    
def start_server():
    try:
        app = make_app()
        app.listen(10001, address="0.0.0.0")
        print("Server is running at http://0.0.0.0:10001")
        IOLoop.current().start()
    except Exception as e:
        print(f"Server error occurred: {e}")
        print("Attempting to restart server in 5 seconds...")
        time.sleep(5)
        restart_server()

def restart_server():
    IOLoop.current().stop()
    lock.release()
    
    print("Restarting server...")
    start_server()

if __name__ == "__main__":    
    git_version = get_git_commit_id()
    print(f'git version:{git_version}')
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_google", action="store_true")
    parser.add_argument("--use_aliyun", action="store_true")
    parser.add_argument("--google_projectid", type=str, help='use project')
    parser.add_argument("--google_topic", type=str, help='use topic')
    parser.add_argument("--google_credentials", type=str, help='use credentials')
    parser.add_argument("--aliyun_topic", type=str, help='use topic')
    parser.add_argument("--use_version", type=str, help='choose version: v1 v2 neurosync all')
    parser.add_argument("--use_envelope", action="store_true")
    args = parser.parse_args()
    use_google = args.use_google
    use_aliyun = args.use_aliyun
    use_envelope = args.use_envelope

    if args.use_version is None:
        use_version = 'v1'
    else:
        use_version = args.use_version
        
    if use_version == 'v1':
        generatorV1 = BSGenerator(ckp=ckp)  
    elif use_version == 'v2':        
        generatorV2 = BSGeneratorV2()   
    elif use_version == 'ns':        
        generatorNS = BSGeneratorNS()   
    elif use_version == 'all':
        generatorV1 = BSGenerator(ckp=ckp)  
        generatorV2 = BSGeneratorV2()  
        generatorNS = BSGeneratorNS()  
        
        
    if use_google:
        
        if args.google_projectid == '':    
            project_id = "yuji-20240819"
        else:
            project_id = args.google_projectid
        
        if args.google_credentials == '':     
            os.environ['GOOGLE_APPLICATION_CREDENTIALS']= os.path.join(os.path.dirname(os.path.abspath(__file__)), 'yuji-20240819-828675ef42a3.json')
        else:
            os.environ['GOOGLE_APPLICATION_CREDENTIALS']= os.path.join(os.path.dirname(os.path.abspath(__file__)), args.google_credentials)
                          
        # topic_id = "projects/yuji-20240819/topics/mouth-data-topic"
        
        if args.google_topic == '':        
            topic_id = "mouth-data-topic"
        else:
            topic_id = args.google_topic
        google_publisher = pubsub_v1.PublisherClient()
        # The `topic_path` method creates a fully qualified identifier
        # in the form `projects/{project_id}/topics/{topic_id}`
        google_topic_path = google_publisher.topic_path(project_id, topic_id)
    if use_aliyun:
            endpoints = "rmq-cn-vc246ac9s16-vpc.cn-hangzhou.rmg,aliyuncs.com:8080"
            credentials = Credentials()
            # if auth enable
            # credentials = Credentials("ak", "sk")
            config = ClientConfiguration(endpoints, credentials)
            
            if args.aliyun_topic == '':        
                aliyun_topic_path = "mouth-test-topic"
            else:
                aliyun_topic_path = args.aliyun_topic
            aliyun_publisher = Producer(config, (aliyun_topic_path,))
        

    
    start_server()
